import copy
import os
import time
import numpy as np
import pandas as pd
import torch
import torch.utils.data as Data
from torchvision.datasets import FashionMNIST
from torchvision import transforms
from torch import nn
from model import LeNet
import matplotlib.pyplot as plt

os.environ["KMP_DUPLICATE_LIB_OK"] = "True"


# 获取数据
def train_val_data():
    data = FashionMNIST(root='./data',
                        train=True,
                        download=True,
                        transform=transforms.Compose([transforms.Resize(size=28), transforms.ToTensor()]))
    # 划分数据集和验证集
    train_data, val_data = Data.random_split(data, [round(len(data)*0.8), round(len(data)*0.2)])
    train_data = Data.DataLoader(train_data,
                                 batch_size=32,
                                 shuffle=True)

    val_data = Data.DataLoader(val_data,
                               batch_size=32,
                               shuffle=True)
    return train_data, val_data


# 训练
def train(model, train_data, val_data, num_epochs, lr=0.001):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # Adam优化器
    optimizer = torch.optim.Adam(model.parameters(), lr)
    # 交叉熵损失
    criterion = nn.CrossEntropyLoss()
    model = model.to(device)
    # 保存模型训练最好的参数
    best_model_wts = copy.deepcopy(model.state_dict())

    # 初始化参数
    best_acc = 0.0
    train_loss_list = []
    val_loss_list = []
    train_acc_list = []
    val_acc_list = []

    for epoch in range(num_epochs):
        start = time.time()
        print("*" * 20)
        print('Epoch {}/{}'.format(epoch+1, num_epochs))

        train_loss = 0.0
        train_acc = 0
        val_loss = 0.0
        val_acc = 0
        train_num = 0
        val_num = 0

        for i, (images, labels) in enumerate(train_data):
            # 获取输入和真值
            images = images.to(device)
            labels = labels.to(device)
            # 训练模式
            model.train()
            outputs = model(images)
            # 根据最大值获取相应标签
            pre_lab = torch.argmax(outputs, dim=1)
            # 计算交叉熵损失
            loss = criterion(outputs, labels)
            # 设置梯度为 0
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            # 梯度优化
            optimizer.step()
            train_loss += loss.item()*images.size(0)
            # 获取训练正确的样本数
            train_acc += torch.sum(pre_lab == labels.data)
            # 累加每轮样本数
            train_num += images.size(0)

        for step, (images, labels) in enumerate(val_data):
            images = images.to(device)
            labels = labels.to(device)
            # 评估模式
            model.eval()
            outputs = model(images)
            pre_lab = torch.argmax(outputs, dim=1)
            loss = criterion(outputs, labels)
            val_loss += loss.item()*images.size(0)
            val_acc += torch.sum(pre_lab ==labels.data)
            val_num += images.size(0)

        train_loss_list.append(train_loss/train_num)
        train_acc_list.append(np.double(train_acc)/train_num)
        val_loss_list.append(np.double(val_loss)/val_num)
        val_acc_list.append(val_acc / val_num)

        print('Train Loss: {:.4f} Train Acc: {:.4f}'.format(train_loss_list[-1], train_acc_list[-1]))
        print('Val Loss: {:.4f} Val Acc: {:.4f}'.format(val_loss_list[-1], val_acc_list[-1]))

        if val_acc_list[-1] > best_acc:
            best_acc = val_acc_list[-1]
            best_model_wts = copy.deepcopy(model.state_dict())

        end = time.time()
        time_elapsed = end - start
        print('Epoch{} 花费时间为: {}m {:.2f}s'.format(epoch+1, time_elapsed//60, time_elapsed%60))

    # 保存最优模型
    torch.save(best_model_wts, f'./checkpoints/LeNet_epoch_{num_epochs}.pth')

    train_loss_list = [loss.cpu().item() if isinstance(loss, torch.Tensor) else loss for loss in train_loss_list]
    val_loss_list = [loss.cpu().item() if isinstance(loss, torch.Tensor) else loss for loss in val_loss_list]
    train_acc_list = [acc.cpu().item() if isinstance(acc, torch.Tensor) else acc for acc in train_acc_list]
    val_acc_list = [acc.cpu().item() if isinstance(acc, torch.Tensor) else acc for acc in val_acc_list]
    # 记录需要绘制图像的参数
    train_process = pd.DataFrame(data={'epoch': range(num_epochs),
                                       'train_loss_list': train_loss_list,
                                       'val_loss_list': val_loss_list,
                                       'train_acc_list': train_acc_list,
                                       'val_acc_list': val_acc_list})

    return train_process


def plot_loss(train_process):
    plt.figure(figsize=(12, 4))
    plt.subplot(1,2,1)
    plt.plot(train_process['epoch'], train_process.train_loss_list, 'ro-', label='train Loss')
    plt.plot(train_process['epoch'], train_process.val_loss_list, 'bs-', label='val loss')
    plt.legend()
    plt.xlabel('epoch')
    plt.ylabel('loss')

    plt.subplot(1,2,2)
    plt.plot(train_process['epoch'], train_process.train_acc_list, 'ro-', label='train acc')
    plt.plot(train_process['epoch'], train_process.val_acc_list, 'bs-', label='val acc')
    plt.legend()
    plt.xlabel('epoch')
    plt.ylabel('acc')
    plt.show()


if __name__ == '__main__':
    # 实例化模型
    LeNet = LeNet()
    train_data, valid_data = train_val_data()
    train_process = train(LeNet, train_data, valid_data, 20)
    plot_loss(train_process)
